package org.hope.nlp.stanford;

import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

import org.hope.nlp.SearchTerm;
import org.hope.nlp.TextParser;
import org.hope.util.Constants;
import org.hope.util.HopeLogger;
import org.hope.util.ReturnData;

import edu.stanford.nlp.ling.CoreAnnotations.NamedEntityTagAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.PartOfSpeechAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.SentencesAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.TextAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.TokensAnnotation;
import edu.stanford.nlp.ling.CoreLabel;
import edu.stanford.nlp.pipeline.Annotation;
import edu.stanford.nlp.pipeline.StanfordCoreNLP;
import edu.stanford.nlp.process.Morphology;
import edu.stanford.nlp.util.CoreMap;

/**
 * @author Thilanka
 */
public class StanfordParser extends TextParser
{

	/** verb, the variations will start with VB */
	private static final String VERB_BASE = "VB";
	/** noun, the variations will start with this */
	private static final String NOUN_BASE = "NN";
	private static final String ADVERB_BASE = "RB";
	private static final String ADJECTIVE_BASE = "JJ";

	private static List<String> requiredWordTypes;

	/** stop words for the context */
	private List<String> stopWords = new ArrayList<String>();

	private StanfordCoreNLP pipeline;

	@Override
	public ReturnData<String> init()
	{
		ReturnData<String> returnData = new ReturnData<String>(Constants.SUCCESS,"");
		
		long t = System.currentTimeMillis();
		Properties props = new Properties();
		props.put( "annotators", "tokenize, ssplit, pos, lemma, ner, parse, dcoref" );
		pipeline = new StanfordCoreNLP( props );

		requiredWordTypes = new ArrayList<String>();
		requiredWordTypes.add( VERB_BASE );
		requiredWordTypes.add( NOUN_BASE );
		requiredWordTypes.add( ADJECTIVE_BASE );
		requiredWordTypes.add( ADVERB_BASE );

		// init the context based stop words
		initStopWords();

		HopeLogger.log( "## Stanford parser initialized in " + ( System.currentTimeMillis() - t ) / 1000 + " seconds" );

		return returnData;
	}

	private void initStopWords()
	{
		// TODO init it from a file
		stopWords.add( "hotel" );
		stopWords.add( "hotels" );
		stopWords.add( "accommodation" );
		stopWords.add( "accommodations" );
		stopWords.add( "area" );
		stopWords.add( "country" );
	}

	@Override
	public List<SearchTerm> getTerms( String text )
	{
		List<SearchTerm> searchTerms = new ArrayList<SearchTerm>();

		// create an empty Annotation just with the given text
		Annotation document = new Annotation( text );

		// run all Annotators on this text
		pipeline.annotate( document );

		// these are all the sentences in this document
		// a CoreMap is essentially a Map that uses class objects as keys and has values with custom types
		List<CoreMap> sentences = document.get( SentencesAnnotation.class );

		for ( CoreMap sentence : sentences )
		{
			// traversing the words in the current sentence
			// a CoreLabel is a CoreMap with additional token-specific methods
			for ( CoreLabel token : sentence.get( TokensAnnotation.class ) )
			{
				SearchTerm searchTerm = new SearchTerm();
				// this is the text of the token
				String word = token.get( TextAnnotation.class );
				// this is the POS tag of the token
				String pos = token.get( PartOfSpeechAnnotation.class );
				// this is the NER label of the token
				String ne = token.get( NamedEntityTagAnnotation.class );

				//TODO remove
				System.out.println( "temp :" + pos + " - " + word );
				if ( requiredWordTypes.contains( getBase( pos ) ) && !stopWords.contains( word ) )
				{
					if(!isBaseForm( pos ))
					{
						word = getBaseWord( word );
					}
					// TODO improve
					word = decorate( word );					
					
					searchTerm.setWord( word );
					searchTerm.setType( ne );
					searchTerm.setPos( pos );
					searchTerms.add( searchTerm );
				}
			}

		}

		return searchTerms;
	}
	
	private String decorate( String word )
	{
		if(word!=null && word.equalsIgnoreCase( "up" ))
		{
			word = "upCountry";
		}
		return word;
	}

	private boolean isBaseForm( String pos )
	{
		return ( pos.length() == 2 );
	}

	private String getBaseWord( String word )
	{
		Morphology morphologyProcess = new Morphology();
		return morphologyProcess.stem( word );
	}
	
	private String getBase( String pos )
	{
		if ( pos.length() > 2 )
		{
			return pos.substring( 0, 2 );
		}

		return pos;
	}
}
